import os
import sys
import argparse
from modules.audit_executor import MultiPlatformAuditExecutor
from config.config import AUDIT_TYPES
import warnings

warnings.filterwarnings("ignore")
from colorama import Fore, Style, init

init()  # 初始化 colorama

banner = """
    ___              _         _____            __  _            __
   /   | ___  ____ _(_)____   / ___/___  ____  / /_(_)___  ___  / /
  / /| |/ _ \/ __ `/ / ___/   \__ \/ _ \/ __ \/ __/ / __ \/ _ \/ / 
 / ___ /  __/ /_/ / (__  )   ___/ /  __/ / / / /_/ / / / /  __/ /  
/_/  |_\___/\__, /_/____/   /____/\___/_/ /_/\__/_/_/ /_/\___/_/   
           /____/                                                  
"""

print(f"{Fore.CYAN}{Style.BRIGHT}{banner}{Style.RESET_ALL}")

# 用户数据存储文件
USER_FILE = "users.txt"


def init_user_file():
    """初始化用户文件（如果不存在则创建）"""
    if not os.path.exists(USER_FILE):
        with open(USER_FILE, "w", encoding="utf-8") as f:
            pass  # 创建空文件


def user_exists(username):
    """检查用户是否已存在"""
    if not os.path.exists(USER_FILE):
        return False

    with open(USER_FILE, "r", encoding="utf-8") as f:
        for line in f:
            parts = line.strip().split(":")
            if len(parts) >= 2 and parts[0] == username:
                return True
    return False


def register_user():
    """注册新用户（用户名+密码，密码可见）"""
    print("\n===== 用户注册 =====")
    username = input("请输入用户名: ").strip()

    if not username:
        print("用户名不能为空！")
        return

    # 检查用户名是否已存在
    if user_exists(username):
        print("该用户名已存在，请重新注册！")
        return

    # 密码输入与确认（可见）
    while True:
        password = input("请输入密码: ").strip()
        if not password:
            print("密码不能为空！")
            continue

        confirm_pwd = input("请确认密码: ").strip()
        if password == confirm_pwd:
            break
        print("两次密码输入不一致，请重新输入！")

    # 保存用户信息（格式：用户名:密码）
    with open(USER_FILE, "a", encoding="utf-8") as f:
        f.write(f"{username}:{password}\n")
    print("注册成功！")


def login_user(username=None, password=None):
    """用户登录验证（密码可见）"""
    if username is None or password is None:
        print("\n===== 用户登录 =====")
        username = input("请输入用户名: ").strip()
        password = input("请输入密码: ").strip()

    if not os.path.exists(USER_FILE):
        print("用户文件不存在，请先注册！")
        return False

    with open(USER_FILE, "r", encoding="utf-8") as f:
        for line in f:
            parts = line.strip().split(":")
            if len(parts) >= 2:
                stored_user, stored_pwd = parts[0], parts[1]
                if stored_user == username and stored_pwd == password:
                    print("登录成功！")
                    return True

    print("用户名或密码错误，登录失败！")
    return False


def get_available_rules(audit_type):
    """获取指定审计类型的可用规则列表"""
    try:
        if audit_type in AUDIT_TYPES:
            return AUDIT_TYPES[audit_type]["rules"]
        else:
            print(f"错误: 不支持的审计类型: {audit_type}")
            return []
    except ImportError:
        print("错误: 无法导入配置模块")
        return []
    except Exception as e:
        print(f"错误: 获取规则列表失败: {e}")
        return []


def parse_rule_input(rule_input, available_rules):
    """解析规则输入，支持数字序号和规则名称"""
    if not rule_input:
        return []

    rules_to_execute = []
    invalid_rules = []

    for item in rule_input.split(','):
        item = item.strip()
        if not item:
            continue

        # 检查是否是数字序号
        if item.isdigit():
            index = int(item) - 1
            if 0 <= index < len(available_rules):
                rule_name = available_rules[index]
                rules_to_execute.append(rule_name)
            else:
                invalid_rules.append(f"序号 {item} (范围: 1-{len(available_rules)})")
        # 检查是否是规则名称
        elif item in available_rules:
            rules_to_execute.append(item)
        else:
            invalid_rules.append(item)

    # 显示无效规则提示
    if invalid_rules:
        print(f"⚠️  以下规则无效: {', '.join(invalid_rules)}")
        print(f"✅ 可用规则: {', '.join(available_rules)}")

    return rules_to_execute


def run_audit(audit_type, description):
    """执行审计任务（包含规则选择功能）"""
    print(f"\n===== 正在执行{description} =====")

    # 获取可用规则列表
    available_rules = get_available_rules(audit_type)
    if not available_rules:
        print("无法获取规则列表，无法执行审计")
        return

    # 询问是否执行特定规则
    specific_rules = None
    use_specific = input("是否执行特定规则? (y/N): ").strip().lower()
    if use_specific == 'y':
        print("\n可用规则:")
        for i, rule in enumerate(available_rules, 1):
            print(f"  {i:2d}. {rule}")

        rule_input = input("请输入规则名称或序号 (多个用逗号分隔): ").strip()
        specific_rules = parse_rule_input(rule_input, available_rules)

        if specific_rules:
            print(f"✅ 将执行以下规则: {', '.join(specific_rules)}")
        else:
            print("❌ 未选择有效规则，将执行全部规则")
            specific_rules = None

    executor = MultiPlatformAuditExecutor()

    try:
        # 执行审计（传入specific_rules参数）
        results = executor.execute_all_rules(audit_type, specific_rules)

        # 生成并显示控制台报告
        executor.print_console_report()

        # 保存JSON报告
        json_report_file = executor.save_report_to_file()
        print(f"\n详细JSON报告已保存到: {json_report_file}")

        # 生成HTML报告
        html_report_file = executor.generate_html_report()
        print(f"HTML可视化报告已保存到: {html_report_file}")

        # 询问是否打开HTML报告
        open_html = input("\n是否立即打开HTML报告? (y/N): ").strip().lower()
        if open_html == 'y':
            import webbrowser
            import os
            webbrowser.open('file://' + os.path.abspath(html_report_file))
            print("正在浏览器中打开HTML报告...")

    except Exception as e:
        print(f"审计执行过程中发生错误: {e}")


def run_audit_directly(audit_type, description, specific_rules_input=None):
    """直接执行审计任务（用于命令行模式）"""
    print(f"\n===== 正在执行{description} =====")

    # 获取可用规则列表
    available_rules = get_available_rules(audit_type)
    if not available_rules:
        print("无法获取规则列表，无法执行审计")
        return

    # 解析规则输入
    specific_rules = None
    if specific_rules_input:
        specific_rules = parse_rule_input(specific_rules_input, available_rules)
        if specific_rules:
            print(f"✅ 将执行以下规则: {', '.join(specific_rules)}")
        else:
            print("❌ 未提供有效规则，将执行全部规则")

    executor = MultiPlatformAuditExecutor()

    try:
        # 执行审计（传入specific_rules参数）
        results = executor.execute_all_rules(audit_type, specific_rules)

        # 生成并显示控制台报告
        executor.print_console_report()

        # 保存JSON报告
        json_report_file = executor.save_report_to_file()
        print(f"\n详细JSON报告已保存到: {json_report_file}")

        # 生成HTML报告
        html_report_file = executor.generate_html_report()
        print(f"HTML可视化报告已保存到: {html_report_file}")

        # 自动打开HTML报告
        import webbrowser
        import os
        webbrowser.open('file://' + os.path.abspath(html_report_file))
        print("正在浏览器中打开HTML报告...")

    except Exception as e:
        print(f"审计执行过程中发生错误: {e}")


def show_second_menu():
    """二级菜单（审计功能选择）"""
    while True:
        print("\n===== 安全配置审计 =====")
        print("1. 操作系统审计")
        print("2. 数据库审计")
        print("3. 应用服务器审计")
        print("0. 返回上级菜单")

        choice = input("请选择操作: ").strip()

        if choice == "1":
            # 操作系统审计
            print("\n----- 操作系统审计 -----")
            print("1. Windows系统审计")
            print("2. Linux系统审计")
            print("0. 返回")

            sys_choice = input("请选择系统类型: ").strip()
            if sys_choice == "1":
                run_audit("windows", "Windows系统安全审计")
            elif sys_choice == "2":
                run_audit("linux", "Linux系统安全审计")
            elif sys_choice == "0":
                continue
            else:
                print("无效选择")

        elif choice == "2":
            # 数据库审计
            print("\n----- 数据库审计 -----")
            print("1. MySQL数据库审计")
            print("2. Redis数据库审计")
            print("0. 返回")

            db_choice = input("请选择数据库类型: ").strip()
            if db_choice == "1":
                run_audit("mysql", "MySQL数据库安全审计")
            elif db_choice == "2":
                run_audit("redis", "Redis数据库安全审计")
            elif db_choice == "0":
                continue
            else:
                print("无效选择")

        elif choice == "3":
            # 应用服务器审计
            print("\n----- 应用服务器审计 -----")
            print("1. Tomcat服务器审计")
            print("2. Nginx服务器审计")
            print("3. Apache服务器审计")
            print("0. 返回")

            app_choice = input("请选择服务器类型: ").strip()
            if app_choice == "1":
                run_audit("tomcat", "Tomcat服务器安全审计")
            elif app_choice == "2":
                run_audit("nginx", "Nginx服务器安全审计")
            elif app_choice == "3":
                run_audit("apache", "Apache服务器安全审计")
            elif app_choice == "0":
                continue
            else:
                print("无效选择")

        elif choice == "0":
            print("返回上级菜单")
            break

        else:
            print("无效选项，请重新选择")


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(
        description="Aegis Sentinel - 多平台安全配置审计系统",
        epilog="""
使用示例:
  python AegisSentinel.py -u admin -p password -c windows                    # Windows系统审计
  python AegisSentinel.py -u admin -p password -d mysql                      # MySQL数据库审计
  python AegisSentinel.py -u admin -p password -w nginx                      # Nginx服务器审计
  python AegisSentinel.py --register -u newuser -p newpass                   # 注册新用户
  python AegisSentinel.py -u admin -p password -c windows -r "1,2"           # 执行序号1,2的规则
  python AegisSentinel.py -u admin -p password -c windows -r "check_firewall,check_updates"  # 执行指定名称规则
  python AegisSentinel.py                                                    # 进入交互式菜单

规则输入说明:
  -r 参数支持两种格式:
    1. 数字序号: "1,2,3" (对应规则列表中的序号)
    2. 规则名称: "check_firewall,check_updates" (完整的规则名称)
        """
    )

    # 登录参数
    parser.add_argument('-u', '--username', help='用户名')
    parser.add_argument('-p', '--password', help='密码')

    # 审计类型参数（互斥组）
    audit_group = parser.add_mutually_exclusive_group()
    audit_group.add_argument('-c', '--os', choices=['windows', 'linux'],
                             help='操作系统审计: windows 或 linux')
    audit_group.add_argument('-d', '--database', choices=['mysql', 'redis'],
                             help='数据库审计: mysql 或 redis')
    audit_group.add_argument('-w', '--web', choices=['tomcat', 'nginx', 'apache'],
                             help='Web服务器审计: tomcat, nginx 或 apache')

    # 规则选择参数
    parser.add_argument('-r', '--rules',
                        help='指定执行的规则，支持数字序号或规则名称，如: "1,2" 或 "check_firewall,check_updates"')

    # 其他参数
    parser.add_argument('--register', action='store_true',
                        help='注册新用户')

    return parser.parse_args()


def handle_command_line(args):
    """处理命令行模式"""
    init_user_file()

    # 注册模式
    if args.register:
        if args.username and args.password:
            # 检查用户是否已存在
            if user_exists(args.username):
                print("错误: 该用户名已存在！")
                return False

            # 注册用户
            with open(USER_FILE, "a", encoding="utf-8") as f:
                f.write(f"{args.username}:{args.password}\n")
            print(f"用户 {args.username} 注册成功！")
            return True
        else:
            print("错误: 注册需要提供用户名和密码")
            return False

    # 审计模式
    if args.username and args.password:
        # 登录验证
        if login_user(args.username, args.password):
            # 执行相应的审计
            if args.os:
                audit_map = {
                    'windows': ('windows', 'Windows系统安全审计'),
                    'linux': ('linux', 'Linux系统安全审计')
                }
                run_audit_directly(audit_map[args.os][0], audit_map[args.os][1], args.rules)
                return True

            elif args.database:
                audit_map = {
                    'mysql': ('mysql', 'MySQL数据库安全审计'),
                    'redis': ('redis', 'Redis数据库安全审计')
                }
                run_audit_directly(audit_map[args.database][0], audit_map[args.database][1], args.rules)
                return True

            elif args.web:
                audit_map = {
                    'tomcat': ('tomcat', 'Tomcat服务器安全审计'),
                    'nginx': ('nginx', 'Nginx服务器安全审计'),
                    'apache': ('apache', 'Apache服务器安全审计')
                }
                run_audit_directly(audit_map[args.web][0], audit_map[args.web][1], args.rules)
                return True
            else:
                print("错误: 请指定审计类型 (-c, -d, 或 -w)")
                return False
        else:
            return False
    else:
        print("错误: 命令行模式需要提供用户名和密码")
        return False


def main():
    """主函数"""
    init_user_file()

    # 解析命令行参数
    args = parse_arguments()

    # 如果有命令行参数，使用命令行模式
    if any([args.username, args.password, args.os, args.database, args.web, args.register, args.rules]):
        handle_command_line(args)
        return

    # 否则进入交互式菜单模式
    print("===== Aegis Sentinel - 安全配置审计系统 =====")

    while True:
        print("\n----- 主菜单 -----")
        print("1. 注册用户")
        print("2. 登录系统")
        print("3. 命令行帮助")
        print("0. 退出系统")

        choice = input("请选择操作: ").strip()
        if choice == "1":
            register_user()

        elif choice == "2":
            if login_user():
                show_second_menu()  # 登录成功后进入二级菜单

        elif choice == "3":
            print("\n----- 命令行使用说明 -----")
            print("使用示例:")
            print("  python AegisSentinel.py -u admin -p password -c windows")
            print("  python AegisSentinel.py -u admin -p password -d mysql")
            print("  python AegisSentinel.py -u admin -p password -w nginx")
            print("  python AegisSentinel.py --register -u newuser -p newpass")
            print("  python AegisSentinel.py -u admin -p password -c windows -r \"1,2\"")
            print("  python AegisSentinel.py -u admin -p password -c windows -r \"check_firewall,check_updates\"")
            print("\n参数说明:")
            print("  -u, --username: 用户名")
            print("  -p, --password: 密码")
            print("  -c, --os: 操作系统审计 (windows/linux)")
            print("  -d, --database: 数据库审计 (mysql/redis)")
            print("  -w, --web: Web服务器审计 (tomcat/nginx/apache)")
            print("  -r, --rules: 指定执行的规则，支持数字序号或规则名称")
            print("  --register: 注册新用户")
            print("\n规则输入说明:")
            print("  -r 参数支持两种格式:")
            print("  1. 数字序号: \"1,2,3\" (对应规则列表中的序号)")
            print("  2. 规则名称: \"check_firewall,check_updates\" (完整的规则名称)")

        elif choice == "0":
            print("感谢使用 Aegis Sentinel，再见！")
            break

        else:
            print("无效选项，请重新选择")


if __name__ == "__main__":
    main()